package com.zyk.leetcode;

import java.util.Arrays;
import java.util.HashSet;

/**
 * @author zhangsan
 * @date 2021/4/22 13:29
 */
public class M0810 {

    // 并查集实现
    public static int[][] floodFill(int[][] image, int sr, int sc, int newColor) {
        int R = image.length, C = image[0].length;
        int area = R * C;       // 0 ~ area
        UnionFind uf = new UnionFind(area);

        int si = index(sr, sc, C);
        for (int r = 0; r < R; r++) {
            for (int c = 0; c < C; c++) {
                if (r > 0 && image[r][c] == image[r - 1][c])
                    uf.union(index(r, c, C), index(r - 1, c, C));
                if (r < R - 1 && image[r][c] == image[r + 1][c])
                    uf.union(index(r, c, C), index(r + 1, c, C));
                if (c > 0 && image[r][c] == image[r][c - 1])
                    uf.union(index(r, c, C), index(r, c - 1, C));
                if (c < C - 1 && image[r][c] == image[r][c + 1])
                    uf.union(index(r, c, C), index(r, c + 1, C));
            }
        }

        for (int r = 0; r < R; r++) {
            for (int c = 0; c < C; c++) {
                if (uf.isSameSet(index(r, c, C), si)) {
                    image[r][c] = newColor;
                }
            }
        }
        return image;
    }

    public static class UnionFind {
        int[] parents;
        // size数组决定把谁关联到谁上边减轻移动过程
        int[] size;
        // help数组是为了加速,找父过程
        int[] help;
        int sets;

        public UnionFind(int n) {
            parents = new int[n];
            for (int i = 0; i < n; i++) {
                parents[i] = i;
            }

            size = new int[n];
            Arrays.fill(size, 1);

            help = new int[n];
            sets = n;
        }

        /**
         * 寻找他所在的集的头节点
         */
        private int findHead(int i) {
            // 重头戏, 查找父的过程要快. 在查找的时候把他们都直接指向头节点. 这样后序查找就能很快地找到
            int helpIndex = 0;
            while (parents[i] != i) {
                help[helpIndex++] = i;  // 记录寻找一路过程中的元素
                i = parents[i];
            }
            // 最后i来到了头元素, 把这里一路的元素直接指向i
            while (helpIndex > 0) {
                parents[help[--helpIndex]] = i;
            }
            return i;       // 返回这个方法要找的东西
        }

        /**
         * 是否同一组
         */
        public boolean isSameSet(int a, int b) {
            return findHead(a) == findHead(b);
        }


        public void union(int i, int j) {
            // 找到他俩的头
            int iHead = findHead(i);
            int jHead = findHead(j);
            if (iHead != jHead) {
                // 不在一起开始合并, 看谁的元素少, 谁移动
                if (size[iHead] >= size[jHead]) {
                    // i的大, 移动j所在的并查集
                    parents[jHead] = iHead;
                    size[iHead] += size[jHead];
                    size[jHead] = 0;
                } else {
                    // j的大, 移动i所在的并查集
                    parents[iHead] = jHead;
                    size[jHead] += size[iHead];
                    size[iHead] = 0;
                }
                sets--; // 移动完了,别忘记把集的数量-1;
            }
        }

        /**
         * 内部现在有多少个集
         */
        public int sets() {
            return sets;
        }
    }



        public static int index(int r, int c, int C) {
            return r * C + c;
        }
        // 递归尝试实现
        public static int[][] floodFill2(int[][] image, int sr, int sc, int newColor) {
            process(image, sr, sc, newColor, image.length, image[0].length, new HashSet<Integer>());
            return image;
        }

        public static void process(int[][] image, int r, int c, int newColor, int R, int C, HashSet<Integer> set) {
            if(r < 0 || r == R || c < 0 || c == C || set.contains(index(r, c, C))) return;
            set.add(index(r, c, C));
            if (r > 0 && image[r][c] == image[r - 1][c])
                process(image, r - 1, c, newColor, R, C, set);
            if (r < R - 1 && image[r][c] == image[r + 1][c])
                process(image, r + 1, c, newColor, R, C, set);
            if (c > 0 && image[r][c] == image[r][c - 1])
                process(image, r, c - 1, newColor, R, C, set);
            if (c < C - 1 && image[r][c] == image[r][c + 1])
                process(image, r, c + 1, newColor, R, C, set);
            image[r][c] = newColor;
        }

    // for test
    public static void main(String[] args) {
        int[][] matrix = {
                {1, 1, 1},
                {1, 1, 0},
                {1, 0, 1},
        };
        int[][] ans = floodFill(matrix, 0, 0, 2);
        for (int[] an : ans) {
            System.out.println(Arrays.toString(an));
        }


        int[][] matrix2 = {
                {1, 1, 1},
                {1, 1, 0},
                {1, 0, 1},
        };
        int[][] ans2 = floodFill2(matrix2, 0, 0, 2);
        for (int[] an : ans2) {
            System.out.println(Arrays.toString(an));
        }
    }

}
